"""
RAG 集成模块
提供医学影像 RAG 检索和 Prompt 增强功能
"""
import os
import sys
import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer, CLIPImageProcessor
import faiss
import numpy as np
import json
from neo4j import GraphDatabase
from typing import Dict, List, Optional, Tuple

# ==================== 配置 ====================
# 项目根目录
PROJECT_ROOT = "/data/Matrix/yy/gjj_pre_merge-v31/KG_RAG_Project"

# Neo4j 配置
NEO4J_URI = os.environ.get("NEO4J_URI", "bolt://localhost:7687")
NEO4J_USER = os.environ.get("NEO4J_USER", "neo4j")
NEO4J_PASSWORD = os.environ.get("NEO4J_PASSWORD", "12345678")

# FAISS 索引文件路径
FAISS_INDEX_PATH = os.path.join(PROJECT_ROOT, "indexes", "image_index.faiss")
ID_MAPPING_PATH = os.path.join(PROJECT_ROOT, "indexes", "id_mapping.json")

# 如果 indexes 目录不存在，尝试根目录
if not os.path.exists(FAISS_INDEX_PATH):
    FAISS_INDEX_PATH = os.path.join(PROJECT_ROOT, "image_index.faiss")
    ID_MAPPING_PATH = os.path.join(PROJECT_ROOT, "id_mapping.json")

# BiomedCLIP 模型路径
BIOMEDCLIP_MODEL_PATH = os.path.join(PROJECT_ROOT, "models", "Biomedclip")
if not os.path.exists(BIOMEDCLIP_MODEL_PATH):
    # 尝试其他可能的路径
    BIOMEDCLIP_MODEL_PATH = os.environ.get("BIOMEDCLIP_MODEL_PATH", "/home/Matrix/Zds/Biomedclip")

VECTOR_DIMENSION = 512

# 全局变量：RAG 检索器实例
_rag_retriever: Optional['RAGRetriever'] = None


class RAGRetriever:
    """RAG 检索器类，提供医学影像检索和知识增强功能"""
    
    def __init__(self):
        """初始化 RAG 检索器"""
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.biomedclip_model = None
        self.biomedclip_image_processor = None
        self.biomedclip_tokenizer = None
        self.faiss_index = None
        self.id_mapping = None
        self.neo4j_driver = None
        
        try:
            self._load_models()
            self._load_knowledge_base()
            print("✅ RAG 检索器初始化成功")
        except Exception as e:
            print(f"⚠️  RAG 检索器初始化失败: {str(e)}")
            import traceback
            traceback.print_exc()
    
    def _load_models(self):
        """加载 BiomedCLIP 模型"""
        if not os.path.exists(BIOMEDCLIP_MODEL_PATH):
            raise FileNotFoundError(f"BiomedCLIP 模型路径不存在: {BIOMEDCLIP_MODEL_PATH}")
        
        print(f"🔄 加载 BiomedCLIP 模型 (路径: {BIOMEDCLIP_MODEL_PATH})...")
        self.biomedclip_model = AutoModel.from_pretrained(BIOMEDCLIP_MODEL_PATH).to(self.device)
        self.biomedclip_image_processor = CLIPImageProcessor.from_pretrained(BIOMEDCLIP_MODEL_PATH)
        self.biomedclip_tokenizer = AutoTokenizer.from_pretrained(BIOMEDCLIP_MODEL_PATH)
        print("✅ BiomedCLIP 模型加载完成")
    
    def _load_knowledge_base(self):
        """加载 FAISS 索引和 Neo4j 连接"""
        # 加载 FAISS 索引
        if os.path.exists(FAISS_INDEX_PATH) and os.path.exists(ID_MAPPING_PATH):
            print(f"🔄 加载 FAISS 索引 (路径: {FAISS_INDEX_PATH})...")
            self.faiss_index = faiss.read_index(FAISS_INDEX_PATH)
            with open(ID_MAPPING_PATH, 'r', encoding='utf-8') as f:
                self.id_mapping = json.load(f)
            
            # 验证索引和映射的一致性
            index_size = self.faiss_index.ntotal
            mapping_size = len(self.id_mapping)
            if index_size != mapping_size:
                print(f"⚠️  警告: FAISS 索引大小 ({index_size}) 与 ID 映射大小 ({mapping_size}) 不匹配")
                print(f"   将使用较小的值: {min(index_size, mapping_size)}")
            else:
                print(f"✅ FAISS 索引加载完成: {index_size} 条记录（索引和映射一致）")
        else:
            print(f"⚠️  FAISS 索引文件不存在: {FAISS_INDEX_PATH}")
            self.faiss_index = None
            self.id_mapping = None
        
        # 连接 Neo4j
        try:
            print(f"🔄 连接 Neo4j (URI: {NEO4J_URI})...")
            self.neo4j_driver = GraphDatabase.driver(
                NEO4J_URI,
                auth=(NEO4J_USER, NEO4J_PASSWORD)
            )
            # 测试连接
            with self.neo4j_driver.session() as session:
                session.run("RETURN 1")
            print("✅ Neo4j 连接成功")
        except Exception as e:
            print(f"⚠️  Neo4j 连接失败: {str(e)}")
            self.neo4j_driver = None
    
    def _get_image_embedding(self, image_path: str) -> Optional[np.ndarray]:
        """提取图像特征向量"""
        try:
            image = Image.open(image_path).convert("RGB")
            inputs = self.biomedclip_image_processor(images=image, return_tensors="pt").to(self.device)
            with torch.no_grad():
                image_features = self.biomedclip_model.get_image_features(**inputs)
                image_features = image_features / image_features.norm(dim=-1, keepdim=True)
            return image_features.cpu().numpy().flatten()
        except Exception as e:
            print(f"❌ 提取图像特征失败 ({image_path}): {str(e)}")
            return None
    
    def _retrieve_similar_cases(self, image_path: str, top_k: int = 3) -> List[str]:
        """检索相似病例 ID"""
        if self.faiss_index is None or self.id_mapping is None:
            print("⚠️  FAISS 索引不可用，无法检索")
            return []
        
        # 提取查询图像特征
        query_embedding = self._get_image_embedding(image_path)
        if query_embedding is None:
            return []
        
        # FAISS 搜索
        query_vector = np.array([query_embedding]).astype('float32')
        distances, indices = self.faiss_index.search(query_vector, top_k)
        
        # 映射到 case_id（添加严格的边界检查）
        similar_case_ids = []
        mapping_len = len(self.id_mapping)
        index_size = self.faiss_index.ntotal if self.faiss_index else 0
        
        # 使用较小的值作为安全边界
        safe_max_index = min(mapping_len, index_size) - 1
        
        for i in indices[0]:
            # 检查索引有效性：必须 >= 0 且 <= 安全最大值
            # FAISS 可能返回 -1 表示无效结果，或者索引可能超出范围
            if 0 <= i <= safe_max_index:
                try:
                    case_id = self.id_mapping[i]["case_id"]
                    similar_case_ids.append(case_id)
                except (KeyError, IndexError, TypeError) as e:
                    print(f"⚠️  索引 {i} 映射失败: {str(e)}")
                    continue
            else:
                if i != -1:  # -1 是 FAISS 的无效结果标记，不需要警告
                    print(f"⚠️  索引 {i} 超出安全范围 [0, {safe_max_index}] (映射长度: {mapping_len}, 索引大小: {index_size})")
        
        # 去重
        similar_case_ids = list(set(similar_case_ids))
        if not similar_case_ids:
            print(f"⚠️  未找到有效的相似病例（所有索引都被过滤）")
        return similar_case_ids
    
    def _extract_knowledge_from_neo4j(self, case_ids: List[str]) -> Dict:
        """从 Neo4j 中提取结构化知识"""
        if self.neo4j_driver is None:
            return {"cases": [], "findings": [], "diseases": []}
        
        knowledge = {
            "cases": [],
            "findings": [],
            "diseases": []
        }
        
        try:
            with self.neo4j_driver.session() as session:
                for case_id in case_ids:
                    case_info = {"case_id": case_id, "report": "", "findings": [], "images": []}
                    
                    # 查询报告和图像
                    query1 = """
                    MATCH (c:MedicalCase {case_id: $case_id})
                    OPTIONAL MATCH (c)-[:HAS_IMAGE]->(img)
                    RETURN c.caption AS caption, collect(img.url) AS urls
                    """
                    result1 = session.run(query1, case_id=case_id).single()
                    if result1:
                        case_info["report"] = result1["caption"] or ""
                        case_info["images"] = result1["urls"] or []
                    
                    # 查询主要发现
                    query2 = """
                    MATCH (c:MedicalCase {case_id: $case_id})-[r:HAS_FINDING]->(f:Finding)
                    RETURN "在 [" + r.on_anatomy + "] 发现 [" + f.name + "]" AS finding_text
                    """
                    result2 = session.run(query2, case_id=case_id)
                    findings = [record["finding_text"] for record in result2]
                    case_info["findings"] = findings
                    knowledge["findings"].extend(findings)
                    
                    knowledge["cases"].append(case_info)
                    
                    # 查询疾病信息
                    query3 = """
                    MATCH (c:MedicalCase {case_id: $case_id})-[:DIAGNOSED_WITH]->(d:Disease)
                    RETURN d.name AS disease_name
                    """
                    result3 = session.run(query3, case_id=case_id)
                    diseases = [record["disease_name"] for record in result3]
                    knowledge["diseases"].extend(diseases)
        
        except Exception as e:
            print(f"⚠️  Neo4j 知识提取失败: {str(e)}")
        
        # 去重
        knowledge["findings"] = list(set(knowledge["findings"]))
        knowledge["diseases"] = list(set(knowledge["diseases"]))
        
        return knowledge
    
    def retrieve_and_enhance(self, image_path: str, original_question: str, top_k: int = 3) -> Tuple[str, Dict]:
        """
        检索相似病例并增强 Prompt
        
        Args:
            image_path: 查询图像路径
            original_question: 原始问题
            top_k: 检索的相似病例数量
        
        Returns:
            (enhanced_prompt, knowledge_dict)
        """
        if self.biomedclip_model is None:
            print("⚠️  BiomedCLIP 模型未加载，返回原始问题")
            return original_question, {"cases": [], "findings": [], "diseases": []}
        
        # 1. 检索相似病例
        similar_case_ids = self._retrieve_similar_cases(image_path, top_k)
        
        if not similar_case_ids:
            print("⚠️  未找到相似病例，返回原始问题")
            return original_question, {"cases": [], "findings": [], "diseases": []}
        
        print(f"🔍 检索到 {len(similar_case_ids)} 个相似病例: {similar_case_ids}")
        
        # 2. 从 Neo4j 提取知识
        knowledge = self._extract_knowledge_from_neo4j(similar_case_ids)
        
        # 3. 构建增强的 Prompt
        enhanced_prompt = self._build_enhanced_prompt(original_question, knowledge)
        
        return enhanced_prompt, knowledge
    
    def _build_enhanced_prompt(self, original_question: str, knowledge: Dict) -> str:
        """构建增强的 Prompt"""
        prompt_parts = [original_question]
        
        if knowledge.get("findings"):
            findings_text = "\n".join([f"- {f}" for f in knowledge["findings"][:5]])
            prompt_parts.append(f"\n参考相似病例的发现：\n{findings_text}")
        
        if knowledge.get("diseases"):
            diseases_text = ", ".join(knowledge["diseases"][:5])
            prompt_parts.append(f"\n相关疾病：{diseases_text}")
        
        if knowledge.get("cases"):
            case_reports = [case["report"] for case in knowledge["cases"][:2] if case.get("report")]
            if case_reports:
                reports_text = "\n".join([f"病例报告：{r}" for r in case_reports])
                prompt_parts.append(f"\n相似病例报告：\n{reports_text}")
        
        return "\n".join(prompt_parts)


def get_rag_retriever() -> RAGRetriever:
    """获取 RAG 检索器实例（单例模式）"""
    global _rag_retriever
    if _rag_retriever is None:
        _rag_retriever = RAGRetriever()
    return _rag_retriever


def is_rag_available() -> bool:
    """检查 RAG 系统是否可用"""
    try:
        retriever = get_rag_retriever()
        return (
            retriever.biomedclip_model is not None and
            retriever.faiss_index is not None and
            retriever.neo4j_driver is not None
        )
    except Exception:
        return False

